import torch
import torch.nn as nn
import torch.nn.functional as F


class maxMarginLoss(nn.Module):

    def __init__(self, cls_num_list, max_m=0.5, weight=None, s=30):
        super(maxMarginLoss, self).__init__()
        m_list = torch.FloatTensor(cls_num_list).cuda()
        self.m_list = m_list
        assert s > 0
        self.s = s
        self.weight = weight

    def forward(self, x, target):
        output = x + 0.1 * torch.log(self.m_list + 1e-7)
        return F.cross_entropy(self.s * output, target, weight=self.weight, reduction="mean")

class maxMarginLoss_kl(nn.Module):
    def __init__(self, cls_num_list, weight=None, s=30):
        super(maxMarginLoss_kl, self).__init__()
        m_list = torch.FloatTensor(cls_num_list).cuda()
        self.m_list = m_list
        assert s > 0
        self.s = s
        self.weight = weight

    def forward(self, x):
        output = x + 0.1 * torch.log(self.m_list + 1e-7)
        return output